package com.zzedu.common.xss;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import org.jsoup.Jsoup;
import org.jsoup.safety.Whitelist;

import java.util.Iterator;


/**
 * XSS工具

 */
public class XSSUtil {

    private static Whitelist whitelist = Whitelist.basic();

    private static String xssWhiteList = "{\"addTags\":[\"p\",\"strong\",\"em\",\"u\",\"s\",\"blockquote\",\"pre\",\"h1\",\"h2\",\"h3\",\"h4\",\"h5\",\"h6\",\"br\",\"ol\",\"li\",\"ul\",\"sub\",\"sup\",\"span\",\"iframe\",\"a\",\"img\"],\"addAttributes\":[{\"tag\":\":all\",\"attributes\":[\"class\",\"id\",\"src\",\"style\"]},{\"tag\":\"pre\",\"attributes\":[\"spellcheck\"]},{\"tag\":\"iframe\",\"attributes\":[\"frameborder\",\"allowfullscreen\",\"src\"]},{\"tag\":\"a\",\"attributes\":[\"rel\",\"target\"]},{\"tag\":\"img\",\"attributes\":[\"align\",\"alt\",\"height\",\"src\",\"title\",\"width\"]}],\"addProtocols\":[{\"tag\":\"xx\",\"attribute\":\"src\",\"protocols\":[\"src\",\"http\",\"https\",\"data\"]}]}";

    //再载入json自定义白名单
    static {
        try {
            JSONObject whiteListJson = JSON.parseObject(xssWhiteList, JSONObject.class);

            //添加标签 addTags
            JSONArray addTagsJsonArr = whiteListJson.getJSONArray("addTags");
            String[] addTagsArr = addTagsJsonArr.toArray(String[]::new);
            whitelist.addTags(addTagsArr);


            //添加属性 addAttributes
            JSONArray addAttrJsonArr = whiteListJson.getJSONArray("addAttributes");
            Iterator<Object> iter = addAttrJsonArr.iterator();
            while (iter.hasNext()) {
                JSONObject attrJsonObj = (JSONObject) iter.next();
                String tag = attrJsonObj.getString("tag");
                JSONArray attrJsonArr = attrJsonObj.getJSONArray("attributes");
                String[] attrArr = attrJsonArr.toArray(String[]::new);
                whitelist.addAttributes(tag, attrArr);
            }


            //添加 addProtocols
            JSONArray addProtoJsonArr = whiteListJson.getJSONArray("addProtocols");
            iter = addProtoJsonArr.iterator();
            while (iter.hasNext()) {
                JSONObject attrJsonObj = (JSONObject) iter.next();
                String tag = attrJsonObj.getString("tag");
                String attribute = attrJsonObj.getString("attribute");
                JSONArray protoJsonArr = attrJsonObj.getJSONArray("protocols");
                String[] protocolArr = protoJsonArr.toArray(String[]::new);
                whitelist.addProtocols(tag, attribute, protocolArr);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private XSSUtil() {

    }

    /**
     * 使用jsoup设置标签放行白名单
     *
     * @param originStr
     * @return
     */
    public static String jsoupCleanRichText(String originStr) {
        return Jsoup.clean(originStr, whitelist);
    }
    /**
     * 使用jsoup设置标签放行白名单
     *
     * @param originStr
     * @return
     */
    public static boolean isValid(String originStr) {
        return Jsoup.isValid(originStr, whitelist);
    }
}
